Add RL sequence-packing and vLLM block-size knobs - #5100
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for sequence packing and custom KV-cache block sizes in RL post-training by adding max_seq_token_per_tpu and vllm_block_size configurations, updating the Tunix adapter to handle packed-sequence segment IDs, and adding corresponding unit tests. The reviewer feedback suggests adding validation to ensure vllm_block_size is a power of 2 and verifying that max_seq_token_per_tpu is at least max_target_length when sequence packing is enabled to prevent runtime failures.
| vllm_block_size: Optional[int] = Field( | ||
| None, | ||
| gt=0, | ||
| description="KV-cache block (page) size for vLLM. None lets the backend pick it from the engine shape.", | ||
| ) |
There was a problem hiding this comment.
vLLM block sizes must be powers of 2 (typically 8, 16, 32, 64, 128, 256). Adding a field validator ensures that any invalid non-power-of-2 block size is caught early during configuration validation rather than causing a startup failure in vLLM.
| vllm_block_size: Optional[int] = Field( | |
| None, | |
| gt=0, | |
| description="KV-cache block (page) size for vLLM. None lets the backend pick it from the engine shape.", | |
| ) | |
| vllm_block_size: Optional[int] = Field( | |
| None, | |
| gt=0, | |
| description="KV-cache block (page) size for vLLM. None lets the backend pick it from the engine shape.", | |
| ) | |
| @field_validator("vllm_block_size") | |
| @classmethod | |
| def validate_vllm_block_size(cls, v: Optional[int]) -> Optional[int]: | |
| if v is not None and (v & (v - 1)) != 0: | |
| raise ValueError("vllm_block_size must be a power of 2.") | |
| return v |
| rollout_vllm_kwargs = { | ||
| "hf_overrides": trainer_config.vllm_hf_overrides, | ||
| "enable_expert_parallel": sampler_config.enable_expert_parallel, | ||
| "enable_prefix_caching": rollout_prefix_caching_enabled(trainer_config), | ||
| # Ensures vLLM model initializes with correct dtype (not float32 default) | ||
| "dtype": trainer_config.weight_dtype.value, | ||
| } | ||
| if trainer_config.vllm_block_size is not None: | ||
| # Pin the KV-cache page size; left unset, the backend derives it from the | ||
| # engine shape, so unrelated engine changes can move it. | ||
| rollout_vllm_kwargs["block_size"] = trainer_config.vllm_block_size |
There was a problem hiding this comment.
When sequence packing is enabled (max_seq_token_per_tpu > 0), a maximal sequence of length max_target_length must fit in a single row. If max_seq_token_per_tpu is configured to be less than max_target_length, Tunix will fail at startup. Adding an early validation check prevents this runtime failure.
if trainer_config.max_seq_token_per_tpu > 0 and trainer_config.max_seq_token_per_tpu < trainer_config.max_target_length:
raise ValueError(
f"max_seq_token_per_tpu ({trainer_config.max_seq_token_per_tpu}) must be greater than or equal to "
f"max_target_length ({trainer_config.max_target_length}) when sequence packing is enabled."
)
rollout_vllm_kwargs = {
"hf_overrides": trainer_config.vllm_hf_overrides,
"enable_expert_parallel": sampler_config.enable_expert_parallel,
"enable_prefix_caching": rollout_prefix_caching_enabled(trainer_config),
# Ensures vLLM model initializes with correct dtype (not float32 default)
"dtype": trainer_config.weight_dtype.value,
}
if trainer_config.vllm_block_size is not None:
# Pin the KV-cache page size; left unset, the backend derives it from the
# engine shape, so unrelated engine changes can move it.
rollout_vllm_kwargs["block_size"] = trainer_config.vllm_block_sizeTunix packs RL training sequences into fixed-token rows when its training config carries `max_seq_token_per_tpu`, and forwards packed segment ids to the model only when the model's call signature has a parameter literally named `segment_ids`. MaxText exposed neither, so packing could not be turned on from an RL config: the CLI rejected the key (not in RLConfig) and the adapter never received the packed ids. - RLConfig gains `max_seq_token_per_tpu` (default 0 = unpacked, one padded row per sequence) and plumbs it into RLTrainingConfig. - TunixMaxTextAdapter accepts `segment_ids` and maps it to MaxText's `decoder_segment_ids`, so packed rows keep per-sequence attention isolation. - RLConfig gains `vllm_block_size` (default None). Unset, tpu_inference derives the KV-cache page size from the engine shape, so unrelated engine changes (max_model_len, TP/DP layout) silently move it; pinning it keeps generation comparable across such experiments. On a Qwen3-0.6B GRPO run (2048 sequences/step, 12288-token rows) packing cut the actor train time from 15.2 s to 6.7 s per step with unchanged rewards.
bb7c509 to
d315757
Compare
1 participant
Description
Tunix packs RL training sequences into fixed-token rows when its training config carries
max_seq_token_per_tpu, and it forwards the packed segment ids to the model only when the model's call signature has a parameter literally namedsegment_ids. MaxText exposed neither, so packing could not be turned on from an RL config: the CLI rejected the key (not inRLConfig), and the adapter would never have received the packed ids.This PR adds:
max_seq_token_per_tputoRLConfig(default0= unpacked, one padded row per sequence, i.e. today's behavior), plumbed intoRLTrainingConfig. When set, a maximal sequence (max_prefill_predict_length+ generation length) must fit in one row; Tunix validates this at startup.segment_idsparameter onTunixMaxTextAdapter.__call__, mapped to MaxText'sdecoder_segment_ids, so packed rows keep per-sequence attention isolation. Without it the adapter falls back to synthesizing a pad mask and packed sequences would attend to each other.vllm_block_sizetoRLConfig(defaultNone, backend-chosen as today). tpu_inference derives the KV-cache page size from the engine shape, so unrelated engine changes (max_model_len, TP/DP layout) silently move it; pinning it keeps generation comparable across such experiments. It is forwarded throughrollout_vllm_kwargsonly when set.Measured on a Qwen3-0.6B GRPO run on TPU v7x (2048 sequences/step, 12288-token rows, prompts ~150 tokens, completions ~4400 tokens): packing cut
actor_train_timefrom 15.2 s to 6.7 s per step, with rewards and completion lengths unchanged. Both defaults leave existing configs untouched.Tests
tests/post_training/unit/tunix_adapter_test.py:segment_idsis forwarded asdecoder_segment_ids, and takes precedence when both are given.tests/post_training/unit/train_rl_test.py:RLConfigdefaults and explicit values for both knobs.Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.